{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 求解二维线性波动方程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 导入模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "using DifferentialEquations, Plots, BenchmarkTools, Sundials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "include(\"geometry.jl\")\n",
    "include(\"operators.jl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 基本配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"./square.neu\""
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "N = 2 # order of approximate polynomials\n",
    "meshFile = \"./square.neu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set element-wise operators\n",
    "Op = Operator(N)\n",
    "# set geomatric informations\n",
    "Geo = Geometry(meshFile, Op);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6×256 Array{Float64,2}:\n",
       " 0.000404645  1.63738e-7   2.36954e-17  …  1.92875e-22  1.08486e-17\n",
       " 5.73909e-5   1.20298e-5   1.94913e-17     6.76036e-20  6.76036e-20\n",
       " 3.72665e-6   0.000404645  1.08486e-17     1.08486e-17  1.92875e-22\n",
       " 0.000727016  3.06546e-6   8.21851e-20     8.21851e-20  1.94913e-17\n",
       " 6.97696e-5   0.000152391  6.76036e-20     1.94913e-17  8.21851e-20\n",
       " 0.000883826  3.88326e-5   1.92875e-22  …  2.36954e-17  2.36954e-17"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = Geo.SPs[:,:,1]\n",
    "y = Geo.SPs[:,:,2]\n",
    "@inbounds u0 = exp.(-1*((x .- 5.0).^2 + y.^2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RHS_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RHS_1 (generic function with 1 method)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function RHS_1(u, p, t)\n",
    "    \"\"\"\n",
    "    Purpose: Evaluate RHS in 2D linear wave equations, discretized on weak form\n",
    "    with a numerical flux\n",
    "    \"\"\"\n",
    "    Op, Geo = p\n",
    "    # set basic informations\n",
    "    K = Geo.K                   # number of elements\n",
    "    Np = Op.Np                  # number of soulution points at each element\n",
    "    Nfp = Geo.Nfp               # number of flux points at each face\n",
    "    Nfaces = Geo.NFaces         # number of faces at each element\n",
    "\n",
    "    # solution points in interior and exterior elements\n",
    "    vmapM = reshape(Geo.vmapM, Nfp*Nfaces, K)\n",
    "    vmapP = reshape(Geo.vmapP, Nfp*Nfaces, K)\n",
    "\n",
    "    # compute weak derivatives\n",
    "    du = zeros(Float64, size(u))\n",
    "\n",
    "    # compute volume contributions\n",
    "    a = 1.; b = 1.\n",
    "    F = a*u; G = b*u\n",
    "    \n",
    "    # calculate divergence of the discontinuous flux part\n",
    "    @inbounds dFdr = Op.Drw*F\n",
    "    @inbounds dFds = Op.Dsw*F\n",
    "    @inbounds dGdr = Op.Drw*G\n",
    "    @inbounds dGds = Op.Dsw*G\n",
    "    @inbounds du = (Geo.rx .* dFdr + Geo.sx .* dFds) + (Geo.ry .* dGdr + Geo.sy .* dGds)\n",
    "\n",
    "    # evaluate '-' and '+' traces of conservative variables\n",
    "    QM = zeros(Float64, Nfp*Nfaces, K)\n",
    "    QP = zeros(Float64, Nfp*Nfaces, K)\n",
    "\n",
    "    # set boundary conditions by modifying positive traces\n",
    "    x = Geo.SPs[:,:,1]; y = Geo.SPs[:,:,2]    \n",
    "    QB = exp.(-1*((x .- t .- 5.0).^2 + (y .- t).^2))\n",
    "    QP = QB[vmapP]\n",
    "\n",
    "    # evaluate primitive variables & flux functions at '-' and '+' traces\n",
    "    fM = a*QM; gM = b*QM\n",
    "    fP = a*QP; gP = b*QP\n",
    "\n",
    "    # compute the Rusonov numerical fluxes\n",
    "    λ = sqrt(a^2 + b^2)\n",
    "\n",
    "    # numerical flux\n",
    "    nx = Geo.Normal[:,:,1]; ny = Geo.Normal[:,:,2]   \n",
    "    nflux = (nx.*(fP + fM) + ny.*(gP + gM) + λ*(QM - QP))/2.0\n",
    "\n",
    "    du = du - Geo.LIFT*(Geo.Fscale.*nflux)\n",
    "    du\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.034808 seconds (5.16 k allocations: 49.001 MiB)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "6×256 Array{Float64,2}:\n",
       "  0.747269   0.0263396     3.52337e-15  …  1.34761e-11  1.23286e-10\n",
       "  0.906147  -0.000511414  -2.03612e-15     2.05156e-11  2.05156e-11\n",
       "  0.748372   0.00130655    5.97617e-15     1.23286e-10  1.34761e-11\n",
       " -0.199664  -0.0105305     1.35125e-15     1.23469e-11  4.67875e-11\n",
       " -0.22831    0.0113226     8.23083e-16     4.67875e-11  1.23469e-11\n",
       "  1.83526    0.0164816     1.49e-16     …  1.68153e-10  1.68153e-10"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob = ODEProblem(RHS_1, u0, (0.0,1.0), (Op, Geo));\n",
    "@time u_1 = solve(prob, Tsit5(), save_everystep=false);\n",
    "u_1[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RHS_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RHS_2 (generic function with 1 method)"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function RHS_2(u, p, t)\n",
    "    \"\"\"\n",
    "    Purpose: Evaluate RHS in 2D linear wave equations, discretized on weak form\n",
    "    with a numerical flux\n",
    "    \"\"\"\n",
    "    Op, Geo = p\n",
    "    # set basic informations\n",
    "    K = Geo.K                   # number of elements\n",
    "    Np = Op.Np                  # number of soulution points at each element\n",
    "    Nfp = Geo.Nfp               # number of flux points at each face\n",
    "    Nfaces = Geo.NFaces         # number of faces at each element\n",
    "\n",
    "    # solution points in interior and exterior elements\n",
    "    vmapM = reshape(Geo.vmapM, Nfp*Nfaces, K)\n",
    "    vmapP = reshape(Geo.vmapP, Nfp*Nfaces, K)\n",
    "\n",
    "    # compute volume contributions\n",
    "    Q = u\n",
    "    a = 1; b = 1\n",
    "    F = a*u; G = b*u\n",
    "\n",
    "    # compute weak derivatives\n",
    "    rhsQ = zeros(Float64, size(Q))\n",
    "\n",
    "    # calculate divergence of the discontinuous flux part\n",
    "\n",
    "    @inbounds dFdr = Op.Drw*F\n",
    "    @inbounds dFds = Op.Dsw*F\n",
    "    @inbounds dGdr = Op.Drw*G \n",
    "    @inbounds dGds = Op.Dsw*G\n",
    "    @inbounds rhsQ = (Geo.rx .* dFdr + Geo.sx .* dFds) + (Geo.ry .* dGdr + Geo.sy .* dGds)\n",
    "\n",
    "\n",
    "    # evaluate '-' and '+' traces of conservative variables\n",
    "    QM = zeros(Float64, Nfp*Nfaces, K)\n",
    "    QP = zeros(Float64, Nfp*Nfaces, K)\n",
    "    @inbounds QM = Q[vmapM]   # get interior value\n",
    "    @inbounds QP = Q[vmapP]   # get exterior value\n",
    "\n",
    "    # set boundary conditions by modifying positive traces\n",
    "    x = Geo.SPs[:,:,1]\n",
    "    y = Geo.SPs[:,:,2]    \n",
    "    QB = exp.(-1*((x .- t .- 5.0).^2 + (y .- t).^2))\n",
    "    QP = QB[vmapP]\n",
    "\n",
    "    # evaluate primitive variables & flux functions at '-' and '+' traces\n",
    "    fM = a*QM\n",
    "    gM = b*QM\n",
    "    fP = a*QP; gP = b*QP\n",
    "\n",
    "    # compute the Rusonov numerical fluxes\n",
    "    λ = sqrt(a^2 + b^2)\n",
    "\n",
    "    # Lifting fluxes\n",
    "    nx = Geo.Normal[:,:,1]\n",
    "    ny = Geo.Normal[:,:,2]\n",
    "    \n",
    "    numerical_flux = nx.*(fP + fM) + ny.*(gP + gM) + λ*(QM - QP)\n",
    "    \n",
    "    rhsQ = rhsQ - Geo.LIFT*(Geo.Fscale.*numerical_flux/2)\n",
    "    rhsQ\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  3.054995 seconds (2.58 M allocations: 202.183 MiB, 2.81% gc time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "6×256 Array{Float64,2}:\n",
       " 0.106176    -4.65882e-8   3.15166e-18  …   1.32419e-11   5.09308e-11\n",
       " 0.0491586    5.43339e-8   9.10395e-18     -2.57533e-12  -2.57533e-12\n",
       " 0.00950412   3.31553e-7   1.72439e-17      5.09308e-11   1.32419e-11\n",
       " 0.173616     2.30125e-8   7.34926e-20     -1.96076e-12   7.07454e-11\n",
       " 0.0597787   -4.1952e-7   -7.69847e-20      7.07454e-11  -1.96076e-12\n",
       " 0.212916    -3.53963e-7   1.4716e-18   …   9.96147e-11   9.96147e-11"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob = ODEProblem(RHS_2, u0, (0.0,1.0), (Op, Geo));\n",
    "@time u_2 = solve(prob, Tsit5(), save_everystep=false);\n",
    "u_2[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6×256 Array{Float64,2}:\n",
       " 0.0990134  1.00604e-12  3.20683e-18  …  1.26642e-14  5.84704e-11\n",
       " 0.0490152  2.57984e-10  9.20705e-18     1.27175e-12  1.27175e-12\n",
       " 0.011109   3.02884e-8   1.78862e-17     5.84704e-11  1.26642e-14\n",
       " 0.177895   6.57401e-11  1.11225e-20     1.54606e-12  1.05052e-10\n",
       " 0.0595873  1.14067e-8   3.19337e-20     1.05052e-10  1.54606e-12\n",
       " 0.216265   2.90669e-9   2.61028e-23  …  1.27711e-10  1.27711e-10"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u = exp.(-1*((x .- 5.0 .- 1.0).^2 + (y .- 1.0).^2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.1183651001667597"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norm(u .- u_2[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RHS_3 (generic function with 1 method)"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function RHS_3(u, p, t)\n",
    "    \"\"\"\n",
    "    Purpose: Evaluate RHS in 2D linear wave equations, discretized on weak form\n",
    "    with a numerical flux\n",
    "    \"\"\"\n",
    "    Op, Geo = p\n",
    "    # set basic informations\n",
    "    K = Geo.K                   # number of elements\n",
    "    Np = Op.Np                  # number of soulution points at each element\n",
    "    Nfp = Geo.Nfp               # number of flux points at each face\n",
    "    Nfaces = Geo.NFaces         # number of faces at each element\n",
    "\n",
    "    # solution points in interior and exterior elements\n",
    "    vmapM = reshape(Geo.vmapM, Nfp*Nfaces, K)\n",
    "    vmapP = reshape(Geo.vmapP, Nfp*Nfaces, K)\n",
    "\n",
    "    # compute volume contributions\n",
    "    a = 1; b = 1\n",
    "    F = a*u; G = b*u\n",
    "\n",
    "    # compute weak derivatives\n",
    "    rhsQ = zeros(Float64, size(u))\n",
    "\n",
    "    # calculate divergence of the discontinuous flux part\n",
    "\n",
    "    @inbounds dFdr = Op.Drw*F\n",
    "    @inbounds dFds = Op.Dsw*F\n",
    "    @inbounds dGdr = Op.Drw*G \n",
    "    @inbounds dGds = Op.Dsw*G\n",
    "    @inbounds rhsQ = (Geo.rx .* dFdr + Geo.sx .* dFds) + (Geo.ry .* dGdr + Geo.sy .* dGds)\n",
    "\n",
    "\n",
    "    # evaluate '-' and '+' traces of conservative variables\n",
    "    QM = zeros(Float64, Nfp*Nfaces, K)\n",
    "    QP = zeros(Float64, Nfp*Nfaces, K)\n",
    "    @inbounds QM = u[vmapM]   # get interior value\n",
    "    @inbounds QP = u[vmapP]   # get exterior value\n",
    "\n",
    "    # set boundary conditions by modifying positive traces\n",
    "    QB = exp.(-1*((Geo.SPs[:,:,1] .- t .- 5.0).^2 + (Geo.SPs[:,:,2] .- t).^2))\n",
    "    QP = QB[vmapP]\n",
    "\n",
    "    # evaluate primitive variables & flux functions at '-' and '+' traces\n",
    "    fM = a*QM; gM = b*QM\n",
    "    fP = a*QP; gP = b*QP\n",
    "\n",
    "    # compute the Rusonov numerical fluxes\n",
    "    nflux = Geo.Normal[:,:,1].*(fP + fM) + Geo.Normal[:,:,2].*(gP + gM) + sqrt(a^2 + b^2)*(QM - QP)\n",
    "    \n",
    "    return rhsQ - Geo.LIFT*(Geo.Fscale.*nflux/2)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 显式格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.183665 seconds (21.89 k allocations: 208.513 MiB, 43.80% gc time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "6×256 Array{Float64,2}:\n",
       " 0.031313   -9.55162e-16   1.5503e-23   …  0.000645866  0.0109112\n",
       " 0.22308     6.80704e-16  -1.64567e-23     0.00268888   0.00268888\n",
       " 0.604347    2.22981e-15   2.78842e-22     0.0109112    0.000645866\n",
       " 0.063987    6.83525e-16   8.98992e-23     0.0033635    0.0171229\n",
       " 0.272228   -3.73189e-15   3.21032e-23     0.0171229    0.0033635\n",
       " 0.0742967  -3.24422e-15   6.35078e-23  …  0.0215739    0.0215739"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob = ODEProblem(RHS_3, u0, (0.0,3.0), (Op, Geo));\n",
    "@time u_3_e = solve(prob, Tsit5(), save_everystep=false);\n",
    "u_3_e[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 隐式格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.168915 seconds (18.15 k allocations: 150.308 MiB, 34.87% gc time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "6×256 Array{Float64,2}:\n",
       " 0.0313665  -3.46377e-14   1.55186e-23  …  0.000644353  0.0109063\n",
       " 0.223027   -7.95449e-15  -1.72048e-23     0.002698     0.002698\n",
       " 0.604225    2.17679e-14   2.67692e-22     0.0109063    0.000644353\n",
       " 0.0640254   3.07273e-14   8.87395e-23     0.00336943   0.0171239\n",
       " 0.272223   -3.79788e-14   2.96424e-23     0.0171239    0.00336943\n",
       " 0.074316   -9.68579e-14   6.41955e-23  …  0.0215696    0.0215696"
      ]
     },
     "execution_count": 120,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob = ODEProblem(RHS_3, u0, (0.0,3.0), (Op, Geo));\n",
    "@time u_3_i = solve(prob,CVODE_BDF(linear_solver=:GMRES),save_everystep=false);\n",
    "u_3_i[2]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.4.1",
   "language": "julia",
   "name": "julia-1.4"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.4.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
